#!/usr/bin/env python
# encoding: utf-8
'''
@author: Justin Ruan
@license: 
@contact: ruanjun@whut.edu.cn
@time: 2019-11-30
@desc:
'''
import os
import pandas as pd

def findall(string, s):
    ret = []
    index = 0
    while True:
        index = string.find(s, index)
        if index != -1:
            ret.append(index)
            index += len(s)
        else:
            break
    return tuple(ret)

def get_project_root():
    path = os.getcwd()
    #约定 项目的代码在src文件夹，其它的目录与它平级
    pos = path.find("src")

    return path[:pos - 1]

def latest_checkpoint(search_dir):

    filename = []
    epoch = []
    loss = []
    accuracy = []

    for ckpt_name in os.listdir(search_dir):
        file_name = os.path.splitext(ckpt_name)[0]
        value = file_name.split("-")
        if len(value) >= 4:
            filename.append(ckpt_name)
            epoch.append(int(value[1]))
            loss.append(float(value[2]))
            accuracy.append(float(value[3]))

    if len(filename) ==  0: return None

    data = {'filename':filename, 'epoch':epoch, 'loss':loss, 'accuracy':accuracy}
    df = pd.DataFrame(data, columns=['filename','epoch','loss','accuracy'])
    result = df.sort_values(['loss', 'accuracy', 'epoch'], ascending=[True, False, False])

    path = "{}/{}".format(search_dir, result.iloc[0,0])
    return path

def latest_checkpoint_loss(search_dir):

    filename = []
    epoch = []
    loss = []
    accuracy = []

    for ckpt_name in os.listdir(search_dir):
        file_name = os.path.splitext(ckpt_name)[0]
        value = file_name.split("-")
        if len(value) >= 3:
            filename.append(ckpt_name)
            epoch.append(int(value[1]))
            loss.append(float(value[2]))

    if len(filename) ==  0: return None

    data = {'filename':filename, 'epoch':epoch, 'loss':loss,}
    df = pd.DataFrame(data, columns=['filename','epoch','loss'])
    result = df.sort_values(['loss', 'epoch'], ascending=[True, False])

    path = "{}/{}".format(search_dir, result.iloc[0,0])
    return path

def clean_checkpoint(search_dir, best_number = 10):
    filename = []
    epoch = []
    loss = []
    accuracy = []

    for ckpt_name in os.listdir(search_dir):
        file_name = os.path.splitext(ckpt_name)[0]
        value = file_name.split("-")
        if len(value) == 4:
            filename.append(ckpt_name)
            epoch.append(int(value[1]))
            loss.append(float(value[2]))
            accuracy.append(float(value[3]))

    if len(filename) <=  best_number: return None

    data = {'filename':filename, 'epoch':epoch, 'loss':loss, 'accuracy':accuracy}
    df = pd.DataFrame(data, columns=['filename','epoch','loss','accuracy'])
    result = df.sort_values(['loss', 'accuracy', 'epoch'], ascending=[True, False, False])

    for ckpt_name in result.iloc[best_number:, 0]:
        path = "{}/{}".format(search_dir, ckpt_name)
        print("deleted", path)
        os.remove(path)